{-# OPTIONS -Wall #-}


module CpsMon.CpsTrafo
  (

    cpsTrans

  ) where


import CpsMon.Lam
import qualified CpsMon.LamM as M

import Control.Applicative


newtype Mon a = Mon { getMon :: Int -> (a,Int) }

instance Functor Mon where
  fmap f ma = Mon $ \s -> let (a,s1) = getMon ma s in (f a, s1)

instance Applicative Mon where
  pure a    = Mon $ \s -> (a,s)
  mf <*> ma = Mon $ \s -> let (f,s1) = getMon mf s
                              (a,s2) = getMon ma s1
                          in (f a, s2)

instance Monad Mon where
  return    = pure
  ma >>= k  = Mon $ \s -> let (a,s1) = getMon ma s in getMon (k a) s1

runMon :: Mon a -> a
runMon ma = fst $ getMon ma 1

freshK :: Mon VarId
freshK = Mon $ \s -> ('k': show s, s+1)

freshVar :: Char -> Mon VarId
freshVar ch = Mon $ \s -> (ch: show s, s+1)


cpsTrans :: M.Expr -> ((Expr -> Expr) -> Expr)
cpsTrans e = runMon (cE e)

cE :: M.Expr -> Mon ((Expr -> Expr) -> Expr)
cE (M.Constant (M.Int i))       = pure $ \k -> k $ (Constant (Int i))
cE (M.Var x)                    = pure $ \k -> k $ Var x
cE (M.Lam x e)                  = do
    newK    <- freshK
    e1      <- cE' e
    return $ \k -> k $ (Lam x (Lam newK (e1 $ (Var newK))))

cE (M.App e0 e1)                = do
    a       <- freshVar 'a'
    e0c     <- cE e0
    e1c     <- cE e1
    return $ \k -> e0c $ (\v0 -> 
                     e1c $ (\v1 -> 
                       (v0 `App` v1) `App` (Lam a (k $ (Var a)))))
             
cE (M.Constant (M.Print))       = do 
    x       <- freshVar 'x'
    newK    <- freshK
    return $ \k -> k $ (Lam x 
                         (Lam newK ((Var newK) `App` (printC `App` (Var x)))))

cE (M.Computation m)            = do
    newK    <- freshK
    mc      <- cM m
    return $ \k -> k $ (Lam newK (mc $ (Var newK)))


printC :: Expr
printC = Constant PrintC


-- underscore @ so translation time application
cE' :: M.Expr -> Mon (Expr -> Expr)
cE' (M.Constant (M.Int i))      = pure $ \k -> k `App` (Constant (Int i))
cE' (M.Var x)                   = pure $ \k -> k `App` Var x
cE' (M.Lam x e)                 = do
    newK    <- freshK
    ec      <- cE' e
    return $ \k -> k `App` (Lam x (Lam newK (ec $ (Var newK))))

cE' (M.App e0 e1)               = do
    e0c     <- cE e0
    e1c     <- cE e1
    return $ \k -> e0c $ (\v0 -> 
                     e1c $ (\v1 -> 
                       (v0 `App` v1) `App` k))
             

cE' (M.Constant (M.Print))      = do 
    x       <- freshVar 'x'
    newK    <- freshK
    return $ \k -> k `App` (Lam x 
                             (Lam newK ((Var newK) `App` (printC `App` (Var x)))))

cE' (M.Computation m)           = do
    newK    <- freshK
    mc      <- cM m
    return $ \k -> k `App` (Lam newK (mc $ (Var newK)))


cM :: M.Stmt -> Mon (Expr -> Expr)
cM (M.Return e)                 = cE' e
cM (M.Seq x e m)                = do 
    ec      <- cE' e
    mc      <- cM m
    return $ \k -> ec $ (Lam x (mc $ k))

cM (M.ExecSeq x e m)            = do 
    v       <- freshVar 'v'
    ec      <- cE' e
    mc      <- cM m
    return $ \k -> ec $ (Lam v ((Var v) `App` (Lam x (mc $ k))))

--
-- Is there an error in the paper with the wrong /time/ 
-- annotation for @lambda v@ ?
--